# -*- coding: utf-8 -*-
"""
全国城镇失业率 vs 预测值（employment_rate）季度相关性分析
结构仿照 CIER 脚本
"""

import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import pearsonr, spearmanr

# ======== 路径配置 ========
PATH_UNEMP = "/Users/linshangjin/25CCM/NKU-C/dataset_csv/每日失业率数据.csv"      # 列：日期, 全国城镇调查失业率(%)
PATH_EMP   = "/Users/linshangjin/25CCM/NKU-C/dataset_csv/data_processed.xlsx"   # 多 sheet：列 date, employment_rate
OUT_DIR    = "/Users/linshangjin/25CCM/NKU-C/t2/失业率/"   # 输出目录
# =========================

os.makedirs(OUT_DIR, exist_ok=True)

# ---- 读全国城镇失业率（CSV）----
df_u = pd.read_csv(PATH_UNEMP, encoding="utf-8")
dcol = "日期" if "日期" in df_u.columns else "date"
df_u[dcol] = pd.to_datetime(df_u[dcol])
col_unemp = "全国城镇调查失业率(%)"
if col_unemp not in df_u.columns:
    raise ValueError(f"失业率文件缺少列：{col_unemp}")
df_u = df_u[[dcol, col_unemp]].rename(columns={dcol:"date", col_unemp:"UNEMP"})
df_u = df_u.set_index("date").sort_index()
unemp_q = df_u["UNEMP"].resample("Q").mean()  # 季均值

# ---- 读预测值（Excel 多 sheet）----
xls = pd.ExcelFile(PATH_EMP)
parts = []
for sh in xls.sheet_names:
    t = xls.parse(sh)
    if {"date","employment_rate"}.issubset(t.columns):
        tmp = t[["date","employment_rate"]].copy()
        tmp["date"] = pd.to_datetime(tmp["date"])
        tmp["employment_rate"] = pd.to_numeric(tmp["employment_rate"], errors="coerce")
        parts.append(tmp)
if not parts:
    raise ValueError("未在 data_processed.xlsx 找到 date & employment_rate")
df_p = pd.concat(parts, ignore_index=True).dropna(subset=["date"])
df_p = df_p.set_index("date").sort_index()
emp_q = df_p["employment_rate"].resample("Q").mean().rename("EMP")

# ---- 合并到季度 ----
df = pd.concat([unemp_q, emp_q], axis=1).dropna()

# ---- 同比 Δ4 ----
def yoy_quarter(s, k=4):
    return s.diff(k)

df["UNEMP_YoY"] = yoy_quarter(df["UNEMP"])
df["EMP_YoY"]   = yoy_quarter(df["EMP"])
df.to_csv(os.path.join(OUT_DIR, "unemp_emp_quarterly.csv"))

# ---- 相关性计算 ----
def corr_pair(a, b):
    idx = a.dropna().index.intersection(b.dropna().index)
    if len(idx) < 6:
        return np.nan, np.nan, np.nan, np.nan, len(idx)
    rP, pP = pearsonr(a.loc[idx], b.loc[idx])
    rS, pS = spearmanr(a.loc[idx], b.loc[idx])
    return rP, pP, rS, pS, len(idx)

rows = []
rP,pP,rS,pS,n = corr_pair(df["UNEMP"], df["EMP"])
rows.append({"type":"levels","pearson":rP,"pearson_p":pP,"spearman":rS,"spearman_p":pS,"n_obs":n})
rP,pP,rS,pS,n = corr_pair(df["UNEMP_YoY"], df["EMP_YoY"])
rows.append({"type":"yoy(Δ4)","pearson":rP,"pearson_p":pP,"spearman":rS,"spearman_p":pS,"n_obs":n})
pd.DataFrame(rows).to_csv(os.path.join(OUT_DIR, "corr_unemp_emp.csv"), index=False)

# ---- 作图 ----
# 叠线
fig = plt.figure()
df[["UNEMP","EMP"]].plot(ax=plt.gca())
plt.title("Unemployment vs EMP (Quarterly Levels)")
plt.tight_layout()
plt.savefig(os.path.join(OUT_DIR, "overlay_levels.png"), dpi=150); plt.close(fig)

fig = plt.figure()
df[["UNEMP_YoY","EMP_YoY"]].dropna().plot(ax=plt.gca())
plt.axhline(0, color="gray", lw=1)
plt.title("Unemployment vs EMP (Quarterly YoY Δ4)")
plt.tight_layout()
plt.savefig(os.path.join(OUT_DIR, "overlay_yoy.png"), dpi=150); plt.close(fig)

# 散点
def scatter_with_fit(x, y, xlabel, ylabel, path):
    fig = plt.figure()
    plt.scatter(x, y, s=30)
    if len(x) >= 3:
        a, b = np.polyfit(x, y, 1)
        grid = np.linspace(min(x), max(x), 100)
        plt.plot(grid, a*grid + b)
    plt.xlabel(xlabel); plt.ylabel(ylabel)
    plt.tight_layout(); plt.savefig(path, dpi=150); plt.close(fig)

scatter_with_fit(df["UNEMP"].dropna(), df["EMP"].dropna(),
                 "Unemployment (level)", "EMP (level)",
                 os.path.join(OUT_DIR, "scatter_levels.png"))
scatter_with_fit(df["UNEMP_YoY"].dropna(), df["EMP_YoY"].dropna(),
                 "Unemployment YoY Δ4", "EMP YoY Δ4",
                 os.path.join(OUT_DIR, "scatter_yoy.png"))

# 滚动相关（4 季窗口）
rollL = df["UNEMP"].rolling(4).corr(df["EMP"])
fig = plt.figure(); rollL.plot(ax=plt.gca())
plt.title("Rolling Corr (w=4): Levels")
plt.tight_layout()
plt.savefig(os.path.join(OUT_DIR, "rolling_corr_levels.png"), dpi=150); plt.close(fig)

rollY = df["UNEMP_YoY"].rolling(4).corr(df["EMP_YoY"])
fig = plt.figure(); rollY.plot(ax=plt.gca())
plt.title("Rolling Corr (w=4): YoY Δ4")
plt.tight_layout()
plt.savefig(os.path.join(OUT_DIR, "rolling_corr_yoy.png"), dpi=150); plt.close(fig)

# 跨相关（±6 季）
def ccf_xy(x, y, maxlag=6):
    x = (x - x.mean())/x.std()
    y = (y - y.mean())/y.std()
    out = {}
    for k in range(-maxlag, maxlag+1):
        if k >= 0:
            xs, ys = x.iloc[:-k or None], y.iloc[k:] if k else y
        else:
            xs, ys = x.iloc[-k:], y.iloc[:k or None]
        if len(xs) > 3:
            out[k] = float(np.corrcoef(xs, ys)[0,1])
    return out

def save_ccf(x, y, fname_csv, fname_png, title):
    c = ccf_xy(x.dropna(), y.dropna(), maxlag=6)
    cdf = pd.DataFrame({"lag_q": list(c.keys()), "corr": list(c.values())}).sort_values("lag_q")
    cdf.to_csv(os.path.join(OUT_DIR, fname_csv), index=False)
    fig = plt.figure()
    plt.bar(cdf["lag_q"], cdf["corr"])
    plt.title(title); plt.xlabel("Lag quarters (k>0: UNEMP leads)")
    plt.tight_layout(); plt.savefig(os.path.join(OUT_DIR, fname_png), dpi=150); plt.close(fig)

save_ccf(df["UNEMP"], df["EMP"], "ccf_levels.csv", "ccf_levels.png",
         "CCF (Levels): UNEMP vs EMP")
save_ccf(df["UNEMP_YoY"], df["EMP_YoY"], "ccf_yoy.csv", "ccf_yoy.png",
         "CCF (YoY Δ4): UNEMP vs EMP")

print("✅ Done. 输出目录：", os.path.abspath(OUT_DIR))
